#include <yolo.h>

YOLO::YOLO(char *model_path, int input_h, int input_w, int classes_num)
{
    // 加载网络模型
	try{
		this->model = cv::dnn::readNet(model_path);
        printf("Init model success!\n");
	}catch(const std::exception){
        printf("Init model failure!\n");
	}

    // 设置检测配置
    this->input_h = input_h;
    this->input_w = input_w;
	
	this->classes_num = classes_num;

	this->feature_num = 3;        // 根据实际情况修改

	int strides[3] = {8, 16, 32}; // 根据实际情况修改
	memcpy(this->strides, strides, sizeof(strides));

	this->batch_size = 1;         // 根据实际情况修改

	this->anchor_num = 3;         // 根据实际情况修改

	float anchors[3][6] = {
		{ 12.0,  16.0,  19.0,  36.0,  40.0,  28.0},
        { 36.0,  75.0,  76.0,  55.0,  72.0, 146.0},
		{142.0, 110.0, 192.0, 243.0, 459.0, 401.0}
	};                            // 根据实际情况修改
	memcpy(this->anchors, anchors, sizeof(anchors));
}

YOLO::~YOLO()
{
    printf("Dein model success!\n");
}

//-------------------------------------------------------------------------------------------------

int YOLO::detect(cv::Mat &image, std::vector<Result> &results, float con_threshold, float nms_threshold)
{
    // 对输入预处理
	// struct timeval tv1, tv2;
	// long t1, t2, time;
    // gettimeofday(&tv1, NULL);

	cv::Mat input = this->preprocess(image);

	// gettimeofday(&tv2, NULL);
    // t1 = tv2.tv_sec - tv1.tv_sec;
    // t2 = tv2.tv_usec - tv1.tv_usec;
    // time = (long)(t1 * 1000 + t2 / 1000);
    // printf("Preprocess times: %dms\n", time);

	// 进行前向传播
    // gettimeofday(&tv1, NULL);

	std::vector<cv::Mat> features = this->forward(input);

	// gettimeofday(&tv2, NULL);
    // t1 = tv2.tv_sec - tv1.tv_sec;
    // t2 = tv2.tv_usec - tv1.tv_usec;
    // time = (long)(t1 * 1000 + t2 / 1000);
    // printf("Forward times: %dms\n", time);

	// 对输出后处理
    // gettimeofday(&tv1, NULL);

	results = this->postprocess(features, con_threshold, nms_threshold);

	// gettimeofday(&tv2, NULL);
    // t1 = tv2.tv_sec - tv1.tv_sec;
    // t2 = tv2.tv_usec - tv1.tv_usec;
    // time = (long)(t1 * 1000 + t2 / 1000);
    // printf("Postprocess times: %dms\n", time);

	// 保存结果图像
	// if(results.size() > 0){
	// 	// 绘制检测结果
	// 	for(int i = 0; i < results.size(); i++){
	// 		// 绘制目标边框
	// 		cv::Rect bbox = cv::Rect(results[i].box_x, results[i].box_y, results[i].box_w, results[i].box_h);
	// 		cv::rectangle(image, bbox, cv::Scalar(0, 255, 0), 2);

	// 		// 绘制目标信息
	// 		char text[256];
	// 		sprintf(text, "label: %d score: %0.2f", results[i].label, results[i].score);
	// 		cv::putText(image, text, bbox.tl(), cv::FONT_HERSHEY_COMPLEX, 1, cv::Scalar(0, 0, 255), 1);
	// 	}

	// 	// 保存结果图像
	// 	cv::imwrite("result.jpg", image);
	// 	printf("Writing result.jpg success!\n");
	// }
	
    return results.size();
}

//-------------------------------------------------------------------------------------------------

cv::Mat YOLO::preprocess(cv::Mat &image)
{
    // 读取原图高宽
    this->image_h = image.rows;
    this->image_w = image.cols;

    // 缩放图像大小
	cv::Mat input; // 模型输入图像
	if(this->image_h == this->input_h && this->image_w == this->input_w){ // 如果输出等于输入
		// 设置缩放参数
		this->rsimg_h = this->input_h; // 设置缩放新图高度
		this->rsimg_w = this->input_w; // 设置缩放新图宽度
		this->rspad_t = 0;             // 缩放上边填充像素
		this->rspad_b = 0;             // 缩放下边填充像素
		this->rspad_l = 0;             // 缩放左边填充像素
		this->rspad_r = 0;             // 缩放右边填充像素

		// 直接返回原图
		input = image;
	}else if(this->image_h != this->image_w){ // 如果原图高宽不等
		// 等比缩放图像
		float hw_ratio = (float) this->image_h / this->image_w; // 原图宽高比例
		if (hw_ratio > 1) { // 原图高大于宽
			// 等比缩放大小
			this->rsimg_h = this->input_h;
			this->rsimg_w = int(this->input_w / hw_ratio); // 注意: 缩放后的宽必须小于缩放前
			cv::resize(image, input, cv::Size(this->rsimg_w, this->rsimg_h), cv::INTER_AREA);
			
			// 填充图像边缘
			this->rspad_t = 0; // 缩放上边填充像素
			this->rspad_b = 0; // 缩放下边填充像素
			this->rspad_l = int((this->input_w - this->rsimg_w) * 0.5);    // 缩放左边填充像素
			this->rspad_r = this->input_w - this->rsimg_w - this->rspad_l; // 缩放右边填充像素
			cv::copyMakeBorder(input, input, this->rspad_t, this->rspad_b, this->rspad_l, this->rspad_r, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
		}else{              // 图像宽大于高
			// 等比缩放大小
			this->rsimg_h = int(this->input_h * hw_ratio); // 注意: 缩放后的高必须小于缩放前
			this->rsimg_w = this->input_w;
			cv::resize(image, input, cv::Size(this->rsimg_w, this->rsimg_h), cv::INTER_AREA);
			
			// 填充图像边缘
			this->rspad_t = int((this->input_h - this->rsimg_h) * 0.5);    // 缩放上边填充像素
			this->rspad_b = this->input_h - this->rsimg_h - this->rspad_t; // 缩放下边填充像素
			this->rspad_l = 0; // 缩放左边填充像素
			this->rspad_r = 0; // 缩放右边填充像素
			cv::copyMakeBorder(input, input, this->rspad_t, this->rspad_b, this->rspad_l, this->rspad_r, cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));
		}
	}else{ // 如果原图高宽相等
		// 设置缩放参数
		this->rsimg_h = this->input_h; // 设置新图缩放高度
		this->rsimg_w = this->input_w; // 设置新图缩放宽度
		this->rspad_t = 0;             // 缩放上边填充像素
		this->rspad_b = 0;             // 缩放下边填充像素
		this->rspad_l = 0;             // 缩放左边填充像素
		this->rspad_r = 0;             // 缩放右边填充像素

		// 直接缩放图像
		cv::resize(image, input, cv::Size(this->rsimg_w, this->rsimg_h), cv::INTER_AREA);
	}
	
	// 保存缩放图像
    // cv::imwrite("input.jpg", input);
    // printf("Writing input image success!\n");

	return input;
}

//-------------------------------------------------------------------------------------------------

std::vector<cv::Mat> YOLO::forward(cv::Mat &input)
{
	// 设置输入张量
	double scalefactor = 1 / 255.0;                         // 像素缩放比例
	cv::Size size = cv::Size(this->input_w, this->input_w); // 模型输入宽高
	cv::Scalar mean = cv::Scalar(0, 0, 0);                  // 对图像减均值
	bool swapRB = true;                                     // 交换红蓝通道
	bool crop = false;                                      // 中心裁剪图像
	int ddepth = CV_32F;                                    // 输出图像位数

	cv::Mat blob = cv::dnn::blobFromImage(input, scalefactor, size, mean, swapRB, crop, ddepth); // 对图像预处理
	this->model.setInput(blob);                                                                  // 设置输入张量

	// 进行前向传播
	std::vector<cv::Mat> features;
	this->model.forward(features, this->model.getUnconnectedOutLayersNames());

	// 输出张量排序
	std::sort(features.begin(), features.end(), [](cv::Mat &A, cv::Mat &B) {return A.size[2] > B.size[2];});

	// 显示张量形状
	// for(int i = 0; i < features.size(); i++){
	// 	printf("output[%d], dim: %d, shape: %dx%dx%dx%dx%d\n", i, features[i].dims, 
	// 		features[i].size[0], features[i].size[1], features[i].size[2], features[i].size[3], features[i].size[4]);
	// }
	// output[0], dim: 5, shape: 1x3x80x80x85
	// output[1], dim: 5, shape: 1x3x40x40x85
	// output[2], dim: 5, shape: 1x3x20x20x85

	return features;
}

//-------------------------------------------------------------------------------------------------

std::vector<Result> YOLO::postprocess(std::vector<cv::Mat> &features, float con_threshold, float nms_threshold)
{
	// 解析输出张量
	std::vector<int> labels;      // 类别标签
	std::vector<cv::Rect> bboxes; // 边框坐标
	std::vector<float> scores;    // 边框得分

	for(int i = 0; i < this->feature_num; i++){ // 遍历输出特征
		this->parse_feature(features[i], this->strides[i], this->anchors[i], con_threshold, labels, bboxes, scores);
	}

	// 非极大值抑制
	std::vector<int> indices; // 结果索引向量
	cv::dnn::NMSBoxes(bboxes, scores, con_threshold, nms_threshold, indices);

	// 设置检测结果
	std::vector<Result> results; 
	for(std::size_t i = 0; i < indices.size(); i++){
		int k = indices[i]; // 结果索引坐标
		Result result;      // 目标检测结果

		result.label = labels[k];
		result.score = scores[k];
		result.box_x = bboxes[k].x;
		result.box_y = bboxes[k].y;
		result.box_w = bboxes[k].width;
		result.box_h = bboxes[k].height;
		results.push_back(result);
		// printf("label: %d, score: %.2f, x: %d, y: %d, w: %d, h: %d\n", labels[k], scores[k], bboxes[k].x, bboxes[k].y, bboxes[k].width, bboxes[k].height);
	}

	return results;
}

void YOLO::parse_feature(cv::Mat feature, int stride, float anchor[], float con_threshold, std::vector<int> &labels, std::vector<cv::Rect> &bboxes, std::vector<float> &scores)
{
	// 设置解析参数
	float *pdata = (float *)feature.data;                 // 特征数据
	int batch_size = this->batch_size;                    // 批次数量
    int anchor_num = this->anchor_num;                    // 锚框数量
	int grid_y_num = (int)ceil((this->input_h / stride)); // 网格高度
	int grid_x_num = (int)ceil((this->input_w / stride)); // 网格宽度
    int output_num = this->classes_num + 5;               // 特征长度:[x,y,w,h,confidence,class_score_0,...,class_score_n]
	
	// 计算缩放比例
	float scale_h = (float)this->image_h / (float)this->rsimg_h; // 缩放高度比例
	float scale_w = (float)this->image_w / (float)this->rsimg_w; // 缩放宽度比例

	// 解析输出特征
	for(int a = 0; a < anchor_num; a++){ // 遍历输出锚框
		float anchor_w = anchor[2 * a];     // 锚框宽度
		float anchor_h = anchor[2 * a + 1]; // 锚框高度

		for(int i = 0; i < grid_y_num; i++){ // 遍历网格高度
			for(int j = 0; j < grid_x_num; j++){ // 遍历网格宽度
				float box_confidence = pdata[4];                // 获取边框置信
				box_confidence = this->sigmoid(box_confidence); // 计算边框置信

				if(box_confidence > con_threshold){ // 边框置信是否大于置信阈值
					for(int k = 5; k < output_num; k++){ // 遍历类别置信
						float cls_confidence = pdata[k];                                 // 获取类别置信
                   		cls_confidence = this->sigmoid(cls_confidence) * box_confidence; // 计算类别置信
						// printf("cls_confidence: %f, box_confidence: %f\n", cls_confidence, box_confidence);

						if(cls_confidence > con_threshold){ // 类别置信是否大于置信阈值
							// 计算边框坐标
							float dx = pdata[0];
							float dy = pdata[1];
							float dw = pdata[2];
							float dh = pdata[3];

							dx = this->sigmoid(dx);
							dy = this->sigmoid(dy);
							dw = this->sigmoid(dw);
							dh = this->sigmoid(dh);

							float x = (dx * 2.0f - 0.5f + (float)j) * stride; // 边框中心x
							float y = (dy * 2.0f - 0.5f + (float)i) * stride; // 边框中心y
							float w = (dw * 2.0f) * (dw * 2.0f) * anchor_w;   // 边框宽度w
							float h = (dh * 2.0f) * (dh * 2.0f) * anchor_h;   // 边框高度y

							// 计算原图坐标
							float x1 = (x - (float)this->rspad_l - 0.5f * w) * scale_w;
							float y1 = (y - (float)this->rspad_t - 0.5f * h) * scale_h;
							float x2 = (x1 + w * scale_w);                       
							float y2 = (y1 + h * scale_h);                       

							int box_x1 = cv::max(0, cv::min(this->image_w - 1, int(x1))); // 坐标从0开始
							int box_y1 = cv::max(0, cv::min(this->image_h - 1, int(y1))); // 坐标从0开始
							int box_x2 = cv::max(0, cv::min(this->image_w - 1, int(x2))); // 坐标从0开始
							int box_y2 = cv::max(0, cv::min(this->image_h - 1, int(y2))); // 坐标从0开始

							// 添加解析结果
							int label = k - 5;                                                              // 类别标签
							cv::Rect bbox = cv::Rect(cv::Point(box_x1, box_y1), cv::Point(box_x2, box_y2)); // 边框坐标
							float score = cls_confidence;                                                   // 类别置信

							labels.push_back(label); // 添加类别标签
							bboxes.push_back(bbox);  // 添加边框坐标
							scores.push_back(score); // 添加类别置信
							// printf("label: %d, score: %.2f, x: %d, y: %d, w: %d, h: %d\n", label, score, bbox.x, bbox.y, bbox.width, bbox.height);
						}
					}
				}
				
				// 移动数据指针
				pdata += output_num;
			}
		}
	}
}

float YOLO::sigmoid(float x)
{
	return (1.0f / ((float)exp((double)(-x)) + 1.0f));
}